import os
import json
from pyecharts import options as opts
from pyecharts.charts import Bar, Pie, Line, Grid, Page
from pyecharts.commons.utils import JsCode
from pyecharts.globals import ThemeType
from pyecharts.components import Table
import datetime
import shutil

class FinanceDataVisualizer:
    def __init__(self, data_dir='finance_data'):
        self.data_dir = data_dir
        self.html_output = os.path.join(os.getcwd(), 'finance_visualization.html')
        self.template_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'template.html')
        
    def load_data(self):
        """Load all data from JSON files"""
        data = {
            'revenue': [],
            'expenditure': [],
            'tax': [],
            'fund': []
        }
        
        # Load data from each directory
        for data_type in data.keys():
            type_dir = os.path.join(self.data_dir, data_type)
            if not os.path.exists(type_dir):
                continue
                
            for filename in os.listdir(type_dir):
                if filename.endswith('.json'):
                    file_path = os.path.join(type_dir, filename)
                    try:
                        with open(file_path, 'r', encoding='utf-8') as f:
                            file_data = json.load(f)
                            data[data_type].extend(file_data)
                    except Exception as e:
                        print(f"Error loading {file_path}: {e}")
        
        return data
    
    def generate_visualization(self):
        """Generate HTML visualization with ECharts"""
        data = self.load_data()
        
        # If we have no data, use the template directly
        if not any(data.values()):
            if os.path.exists(self.template_path):
                shutil.copy(self.template_path, self.html_output)
                print(f"No data found. Using template directly: {self.html_output}")
                return
        
        # If we do have data, we would update the template with real data
        # For now, just use the template as is
        if os.path.exists(self.template_path):
            shutil.copy(self.template_path, self.html_output)
            print(f"Template-based visualization generated at: {self.html_output}")
        else:
            # If template doesn't exist, generate a basic visualization with pyecharts
            self._generate_basic_visualization(data)
    
    def _generate_basic_visualization(self, data):
        """Generate a basic visualization using pyecharts"""
        # Create a page to hold multiple charts
        page = Page(layout=Page.DraggablePageLayout, page_title="财政部财政收支数据可视化")
        
        # Add title and source information
        title = "财政部财政收支数据可视化"
        source = "数据来源：中华人民共和国财政部官方网站"
        
        # Create fiscal overview section
        fiscal_overview = self.create_fiscal_overview(data)
        
        # Create revenue analysis section
        revenue_analysis = self.create_revenue_analysis(data)
        
        # Create expenditure analysis section
        expenditure_analysis = self.create_expenditure_analysis(data)
        
        # Create government fund section
        fund_analysis = self.create_fund_analysis(data)
        
        # Create original data table
        data_table = self.create_data_table(data)
        
        # Add components to page
        page.add(fiscal_overview)
        page.add(revenue_analysis)
        page.add(expenditure_analysis)
        page.add(fund_analysis)
        page.add(data_table)
        
        # Generate HTML file
        page.render(self.html_output)
        print(f"Basic visualization generated at: {self.html_output}")
        
    def create_fiscal_overview(self, data):
        """Create fiscal overview section"""
        # Get the latest revenue and expenditure data
        latest_revenue = None
        latest_revenue_amount = 0
        latest_expenditure = None
        latest_expenditure_amount = 0
        
        # Find latest national revenue
        for item in data['revenue']:
            if item.get('level') == 'national' and item.get('amount') is not None:
                if latest_revenue is None or (item.get('year'), item.get('month')) > (latest_revenue.get('year'), latest_revenue.get('month')):
                    latest_revenue = item
                    latest_revenue_amount = item.get('amount')
        
        # Find latest national expenditure
        for item in data['expenditure']:
            if item.get('level') == 'national' and item.get('amount') is not None:
                if latest_expenditure is None or (item.get('year'), item.get('month')) > (latest_expenditure.get('year'), latest_expenditure.get('month')):
                    latest_expenditure = item
                    latest_expenditure_amount = item.get('amount')
        
        # Create the chart
        grid = Grid()
        
        # Create title for the section
        title = "财政收支概览"
        
        # Create bar chart for fiscal overview
        bar = (
            Bar()
            .add_xaxis(["财政收入", "财政支出"])
            .add_yaxis("金额(亿元)", [latest_revenue_amount, latest_expenditure_amount])
            .set_global_opts(
                title_opts=opts.TitleOpts(title=title),
                toolbox_opts=opts.ToolboxOpts()
            )
        )
        
        return bar
        
    def create_revenue_analysis(self, data):
        """Create revenue analysis section"""
        # Filter revenue data
        revenue_data = data['revenue']
        
        # Create charts for revenue analysis
        pie = (
            Pie()
            .add(
                "收入构成",
                [("税收收入", 70), ("非税收入", 30)],  # Sample data, replace with actual data
                radius=["40%", "75%"],
            )
            .set_global_opts(
                title_opts=opts.TitleOpts(title="财政收入分析"),
                legend_opts=opts.LegendOpts(orient="vertical", pos_top="15%", pos_left="2%")
            )
        )
        
        return pie
        
    def create_expenditure_analysis(self, data):
        """Create expenditure analysis section"""
        # Create expenditure pie charts
        pie1 = (
            Pie()
            .add(
                "支出构成",
                [
                    ("教科文卫", 23), 
                    ("社会保障", 18),
                    ("一般公共服务", 14),
                    ("国防", 10),
                    ("农林水事务", 8),
                    ("城乡社区", 6),
                    ("公共安全", 6),
                    ("交通运输", 5),
                    ("其他", 10)
                ],  # Sample data, replace with actual data
                radius=["0%", "70%"],
            )
            .set_global_opts(
                title_opts=opts.TitleOpts(title="财政支出分析"),
                legend_opts=opts.LegendOpts(orient="vertical", pos_top="15%", pos_left="2%")
            )
        )
        
        return pie1
        
    def create_fund_analysis(self, data):
        """Create government fund section"""
        # Create government fund chart
        bar = (
            Bar()
            .add_xaxis(["政府性基金收入", "政府性基金支出"])
            .add_yaxis("", [5000, 8000])  # Sample data, replace with actual data
            .set_global_opts(
                title_opts=opts.TitleOpts(title="政府性基金预算"),
                toolbox_opts=opts.ToolboxOpts()
            )
        )
        
        return bar
        
    def create_data_table(self, data):
        """Create original data table"""
        # Sample data for the table
        sample_data = [
            ["一般公共预算收入", "60877", "-2.3"],
            ["税收收入", "49372", "-8.4"],
            ["非税收入", "11705", "+10.1"],
            ["增值税", "29046", "-7.1"]
        ]
        
        # Create table
        table = Table()
        headers = ["收入项目", "金额 (亿元)", "同比增长 (%)"]
        table.add(headers, sample_data).set_global_opts(
            title_opts=opts.ComponentTitleOpts(title="原始数据")
        )
        
        return table

# If run directly, generate visualization
if __name__ == "__main__":
    visualizer = FinanceDataVisualizer()
    visualizer.generate_visualization() 